{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# 使用Python代码准备训练样本",
   "id": "4fb0b5720010c433"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "import matplotlib.pyplot as plt\n",
    "# 设置中文显示\n",
    "plt.rcParams['font.sans-serif'] = ['SimHei']\n",
    "plt.rcParams['axes.unicode_minus'] = False"
   ],
   "id": "2ff3e9842a2efd59",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## 1. 方法一：基于现有分类结果的半自动样本生成",
   "id": "62b2e0e4d01d2f22"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "import numpy as np\n",
    "import rasterio\n",
    "import geopandas as gpd\n",
    "from shapely.geometry import Point\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.cluster import KMeans\n",
    "import cv2\n",
    "import os\n",
    "\n",
    "def generate_training_samples(image_path, output_shp, num_samples_per_class=50):\n",
    "    \"\"\"\n",
    "    半自动生成训练样本\n",
    "    :param image_path: 输入遥感影像路径\n",
    "    :param output_shp: 输出样本shapefile路径\n",
    "    :param num_samples_per_class: 每类样本数量\n",
    "    \"\"\"\n",
    "    # 1. 加载影像数据\n",
    "    with rasterio.open(image_path) as src:\n",
    "        img = src.read()\n",
    "        profile = src.profile\n",
    "        bounds = src.bounds\n",
    "\n",
    "    # 转置为(height, width, bands)格式\n",
    "    img = np.transpose(img, (1, 2, 0))\n",
    "    height, width, bands = img.shape\n",
    "\n",
    "    # 2. 预处理影像\n",
    "    # 归一化到0-1范围\n",
    "    img_norm = (img - np.min(img)) / (np.max(img) - np.min(img))\n",
    "    # 转换为8bit用于显示\n",
    "    img_display = (img_norm * 255).astype(np.uint8)\n",
    "    # 只取前3个波段生成RGB图像\n",
    "    rgb = img_display[:, :, :3]\n",
    "\n",
    "    # 3. 使用K-means进行预分类\n",
    "    print(\"正在进行预分类...\")\n",
    "    pixels = img_norm.reshape(-1, bands)\n",
    "    kmeans = KMeans(n_clusters=8, random_state=42).fit(pixels)\n",
    "    labels = kmeans.labels_.reshape(height, width)\n",
    "\n",
    "    # 4. 从每个类别中随机选择样本点\n",
    "    samples = []\n",
    "    for class_id in range(8):\n",
    "        # 获取当前类别的所有像素坐标\n",
    "        y, x = np.where(labels == class_id)\n",
    "        if len(y) == 0:\n",
    "            continue\n",
    "\n",
    "        # 随机选择样本点\n",
    "        indices = np.random.choice(len(y), size=min(num_samples_per_class, len(y)), replace=False)\n",
    "\n",
    "        for idx in indices:\n",
    "            # 转换为地理坐标\n",
    "            px, py = x[idx], y[idx]\n",
    "            lon = bounds.left + (px / width) * (bounds.right - bounds.left)\n",
    "            lat = bounds.top - (py / height) * (bounds.top - bounds.bottom)\n",
    "\n",
    "            samples.append({\n",
    "                'class': class_id + 1,  # 类别ID从1开始\n",
    "                'geometry': Point(lon, lat),\n",
    "                'pixel_x': px,\n",
    "                'pixel_y': py\n",
    "            })\n",
    "\n",
    "    # 5. 保存为Shapefile\n",
    "    gdf = gpd.GeoDataFrame(samples, crs=profile['crs'])\n",
    "    gdf.to_file(output_shp)\n",
    "    print(f\"已生成{len(samples)}个样本点，保存至: {output_shp}\")\n",
    "\n",
    "    return gdf\n",
    "\n",
    "# 使用示例\n",
    "image_path = \"../data/output_preprocessed/preprocessed_landsat8_2021.tif\"\n",
    "output_shp = \"../data/training_samples/auto_samples/auto_samples_xinDu_2021.shp\"\n",
    "samples_gdf = generate_training_samples(image_path, output_shp, num_samples_per_class=50)"
   ],
   "id": "bb7acc8ee7eb1cc",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "samples_gdf",
   "id": "866424a6c1356e07",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### 可视化查看",
   "id": "ec23e3ae4151bbd4"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import geopandas as gpd\n",
    "import rasterio\n",
    "from rasterio.plot import show\n",
    "\n",
    "# 1. 加载样本点和原始影像\n",
    "sample_shp = \"../data/training_samples/auto_samples/auto_samples.shp\"\n",
    "image_path = \"../data/output_preprocessed/preprocessed_landsat8_2021.tif\"\n",
    "\n",
    "samples = gpd.read_file(sample_shp)\n",
    "with rasterio.open(image_path) as src:\n",
    "    img = src.read([4, 3, 2])  # 假彩色合成（近红外-红-绿）\n",
    "    transform = src.transform\n",
    "\n",
    "# 2. 创建可视化画布\n",
    "fig, ax = plt.subplots(figsize=(12, 10))\n",
    "\n",
    "# 3. 显示遥感影像（假彩色）\n",
    "show(img, ax=ax, transform=transform, cmap='viridis')\n",
    "\n",
    "# 4. 叠加样本点（按类别着色）\n",
    "classes = samples['class'].unique()\n",
    "colors = plt.cm.tab10(range(len(classes)))  # 使用10种颜色\n",
    "\n",
    "for cls, color in zip(classes, colors):\n",
    "    cls_samples = samples[samples['class'] == cls]\n",
    "    cls_samples.plot(ax=ax, color=color, markersize=50, label=f'Class {cls}')\n",
    "\n",
    "# 5. 添加图例和标题\n",
    "ax.legend(title=\"样本类别\")\n",
    "ax.set_title(\"遥感影像分类样本点分布\", fontsize=14)\n",
    "plt.tight_layout()\n",
    "\n",
    "# 6. 保存或显示\n",
    "# plt.savefig(\"../output/samples_visualization.png\", dpi=300, bbox_inches='tight')\n",
    "plt.show()"
   ],
   "id": "890b82c46594bca9",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20, 10))\n",
    "\n",
    "# 左图：仅原始影像\n",
    "show(img, ax=ax1, transform=transform, title=\"原始影像（假彩色）\")\n",
    "\n",
    "# 右图：影像+样本点\n",
    "show(img, ax=ax2, transform=transform)\n",
    "samples.plot(ax=ax2, column='class', legend=True, markersize=30,\n",
    "             legend_kwds={'title': '样本类别'})\n",
    "ax2.set_title(\"样本点分布\")\n",
    "\n",
    "# plt.savefig(\"../output/samples_comparison.png\", dpi=300)\n",
    "plt.show()"
   ],
   "id": "32361e235740ea0e",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## 方法二：交互式样本标注工具",
   "id": "402d15ddf26ecd97"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "import numpy as np\n",
    "import rasterio\n",
    "import geopandas as gpd\n",
    "from shapely.geometry import Point\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "\n",
    "class InteractiveSampleLabeler:\n",
    "    def __init__(self, image_path):\n",
    "        self.image_path = image_path\n",
    "        self.samples = []\n",
    "        self.current_class = 1\n",
    "        self.class_colors = {\n",
    "            1: (0, 255, 0),    # 绿色 - 林地\n",
    "            2: (0, 255, 255),  # 黄色 - 草地\n",
    "            3: (0, 0, 255),    # 红色 - 城镇\n",
    "            4: (255, 255, 0),  # 青色 - 水域\n",
    "            5: (0, 128, 0),    # 深绿 - 灌木\n",
    "            6: (255, 165, 0),  # 橙色 - 农田\n",
    "            7: (128, 128, 128),# 灰色 - 裸地\n",
    "            8: (255, 255, 255) # 白色 - 其他\n",
    "        }\n",
    "\n",
    "        # 加载影像\n",
    "        with rasterio.open(image_path) as src:\n",
    "            self.img = src.read()\n",
    "            self.profile = src.profile\n",
    "            self.bounds = src.bounds\n",
    "\n",
    "        # 准备显示图像\n",
    "        self.rgb = np.transpose(self.img[:3], (1, 2, 0))\n",
    "        self.rgb = cv2.normalize(self.rgb, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)\n",
    "        self.display_img = self.rgb.copy()\n",
    "\n",
    "        # 创建窗口\n",
    "        cv2.namedWindow('Sample Labeler')\n",
    "        cv2.setMouseCallback('Sample Labeler', self.mouse_callback)\n",
    "\n",
    "        print(\"使用说明:\")\n",
    "        print(\"1. 按键1-8选择当前类别\")\n",
    "        print(\"2. 鼠标左键点击添加样本\")\n",
    "        print(\"3. 按's'保存样本\")\n",
    "        print(\"4. 按'q'退出\")\n",
    "\n",
    "    def mouse_callback(self, event, x, y, flags, param):\n",
    "        if event == cv2.EVENT_LBUTTONDOWN:\n",
    "            # 添加样本点\n",
    "            lon = self.bounds.left + (x / self.rgb.shape[1]) * (self.bounds.right - self.bounds.left)\n",
    "            lat = self.bounds.top - (y / self.rgb.shape[0]) * (self.bounds.top - self.bounds.bottom)\n",
    "\n",
    "            self.samples.append({\n",
    "                'class': self.current_class,\n",
    "                'geometry': Point(lon, lat),\n",
    "                'pixel_x': x,\n",
    "                'pixel_y': y\n",
    "            })\n",
    "\n",
    "            # 在图像上标记\n",
    "            cv2.circle(self.display_img, (x, y), 5, self.class_colors[self.current_class], -1)\n",
    "            cv2.imshow('Sample Labeler', self.display_img)\n",
    "\n",
    "    def run(self):\n",
    "        while True:\n",
    "            cv2.imshow('Sample Labeler', self.display_img)\n",
    "            key = cv2.waitKey(1) & 0xFF\n",
    "\n",
    "            # 设置当前类别\n",
    "            if ord('1') <= key <= ord('8'):\n",
    "                self.current_class = key - ord('0')\n",
    "                print(f\"当前类别: {self.current_class}\")\n",
    "\n",
    "            # 保存样本\n",
    "            elif key == ord('s'):\n",
    "                if len(self.samples) > 0:\n",
    "                    gdf = gpd.GeoDataFrame(self.samples, crs=self.profile['crs'])\n",
    "                    output_path = os.path.splitext(self.image_path)[0] + \"_samples.shp\"\n",
    "                    gdf.to_file(output_path)\n",
    "                    print(f\"已保存{len(self.samples)}个样本到: {output_path}\")\n",
    "                else:\n",
    "                    print(\"没有样本可保存\")\n",
    "\n",
    "            # 退出\n",
    "            elif key == ord('q'):\n",
    "                break\n",
    "\n",
    "        cv2.destroyAllWindows()\n",
    "\n",
    "# 使用示例\n",
    "labeler = InteractiveSampleLabeler(\"../data/preprocessed_output/preprocessed_landsat8_2021.tif\")\n",
    "labeler.run()"
   ],
   "id": "2c684978eb54ccc5",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "d7eb787578c896d5",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "e5ca6326bad81ac4",
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
